"""
数据库迁移脚本 - 添加UserSession刷新令牌字段
"""
import sys
import os
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from sqlalchemy import text
from core.database import SessionLocal, engine
from core.logging_config import get_logger

logger = get_logger("db_migration")


def migrate_database():
    """执行数据库迁移"""
    logger.info("开始执行数据库迁移...")
    
    # 检查是否需要添加refresh_token字段
    with engine.connect() as conn:
        result = conn.execute(text("PRAGMA table_info(user_session)"))
        columns = [row[1] for row in result.fetchall()]
        
        if 'refresh_token' not in columns:
            logger.info("添加refresh_token字段到user_session表...")
            conn.execute(text("ALTER TABLE user_session ADD COLUMN refresh_token VARCHAR(255)"))
            conn.execute(text("CREATE INDEX idx_user_session_refresh_token ON user_session(refresh_token)"))
            logger.info("refresh_token字段添加完成")
        
        if 'refresh_token_expires_at' not in columns:
            logger.info("添加refresh_token_expires_at字段到user_session表...")
            conn.execute(text("ALTER TABLE user_session ADD COLUMN refresh_token_expires_at DATETIME"))
            logger.info("refresh_token_expires_at字段添加完成")
        
        conn.commit()
    
    logger.info("数据库迁移完成")


if __name__ == "__main__":
    migrate_database()